# backend/utils/websocket_test.py
import asyncio
import websockets
import json
from typing import Dict, Any


class WebSocketTestClient:
    """WebSocket测试客户端"""
    
    def __init__(self, uri: str, token: str):
        self.uri = uri
        self.token = token
        self.websocket = None
        self.messages = []
    
    async def connect(self):
        """连接WebSocket"""
        try:
            self.websocket = await websockets.connect(f"{self.uri}?token={self.token}")
            print(f"已连接到 {self.uri}")
        except Exception as e:
            print(f"连接失败: {str(e)}")
            return False
        return True
    
    async def disconnect(self):
        """断开连接"""
        if self.websocket:
            await self.websocket.close()
            print("连接已断开")
    
    async def send_message(self, message_type: str, data: Dict[str, Any]):
        """发送消息"""
        if not self.websocket:
            print("未连接")
            return
        
        message = {
            "type": message_type,
            "data": data
        }
        
        await self.websocket.send(json.dumps(message))
        print(f"已发送: {message}")
    
    async def receive_messages(self):
        """接收消息"""
        if not self.websocket:
            return
        
        try:
            async for message in self.websocket:
                data = json.loads(message)
                self.messages.append(data)
                print(f"收到消息: {data}")
        except websockets.exceptions.ConnectionClosed:
            print("连接已关闭")
        except Exception as e:
            print(f"接收消息错误: {str(e)}")
    
    async def start_listening(self):
        """开始监听消息"""
        asyncio.create_task(self.receive_messages())
    
    def get_messages(self, message_type: str = None):
        """获取消息"""
        if message_type:
            return [msg for msg in self.messages if msg.get("type") == message_type]
        return self.messages
    
    def clear_messages(self):
        """清空消息"""
        self.messages.clear()


# 测试用例示例

async def test_websocket_chat():
    """测试WebSocket聊天功能"""
    
    client = WebSocketTestClient("ws://localhost:8000/ws", "test_token")
    
    # 连接
    if not await client.connect():
        return
    
    # 开始监听
    await client.start_listening()
    
    # 发送聊天消息
    await client.send_message("chat_message", {
        "message": "Hello, this is a test message",
        "project_id": "test_project"
    })
    
    # 等待响应
    await asyncio.sleep(2)
    
    # 检查收到的消息
    messages = client.get_messages("chat_response")
    print(f"收到 {len(messages)} 条聊天回复")
    
    # 断开连接
    await client.disconnect()


async def test_websocket_code_collaboration():
    """测试WebSocket代码协作功能"""
    
    client1 = WebSocketTestClient("ws://localhost:8000/ws", "user1_token")
    client2 = WebSocketTestClient("ws://localhost:8000/ws", "user2_token")
    
    # 两个客户端连接
    await client1.connect()
    await client2.connect()
    
    # 开始监听
    await client1.start_listening()
    await client2.start_listening()
    
    # 客户端1发送代码变更
    await client1.send_message("code_change", {
        "file_id": "test_file",
        "project_id": "test_project",
        "changes": [
            {"type": "insert", "line": 1, "text": "print('Hello World')"}
        ]
    })
    
    # 等待
    await asyncio.sleep(1)
    
    # 检查客户端2是否收到变更
    code_changes = client2.get_messages("code_change")
    print(f"客户端2收到 {len(code_changes)} 条代码变更")
    
    # 断开连接
    await client1.disconnect()
    await client2.disconnect()


if __name__ == "__main__":
    # 运行测试
    asyncio.run(test_websocket_chat())
    asyncio.run(test_websocket_code_collaboration())
